import re
import ast
import inspect
from enum import Enum
from typing import Any, Dict, List, Optional, Type, get_type_hints, Literal, Tuple


def _get_type_str(node: ast.Subscript) -> str:
    """处理泛型类型如List[int]的字符串表示"""
    base = node.value.id if isinstance(node.value, ast.Name) else "any"  # 类型名称（如List）
    slice_str = _get_ast_type_info(node.slice).get("type", "any")  # 泛型参数
    return f"{base}[{slice_str}]"


def _get_ast_type_info(node: ast.expr) -> Dict[str, Any]:
    """
    从 AST 节点中获取类型信息（改进版）
    """

    # 处理负数情况
    def _get_constant_value(node: ast.expr):
        if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub):
            if isinstance(node.operand, ast.Constant):
                return -node.operand.value
        elif isinstance(node, ast.Constant):
            return node.value
        return None

    # 修改联合类型处理逻辑
    if isinstance(node, ast.BinOp) and isinstance(node.op, ast.BitOr):

        def _collect_types(n):
            """递归收集类型名称"""
            if isinstance(n, ast.BinOp) and isinstance(n.op, ast.BitOr):
                return _collect_types(n.left) + _collect_types(n.right)
            if isinstance(n, ast.Name):
                return [n.id]
            if isinstance(n, ast.Subscript):
                return [_get_type_str(n)]  # 处理泛型类型
            return [_get_ast_type_info(n).get("type", "any")]  # 兜底处理

        types = _collect_types(node)
        return {"type": " | ".join(sorted(set(types)))}

    if isinstance(node, ast.Name):
        return {"type": node.id}
    elif isinstance(node, ast.Subscript):
        if isinstance(node.value, ast.Name):
            # 处理泛型类型
            if node.value.id == "List":
                return {"type": "array", "items": _get_ast_type_info(node.slice)}
            elif node.value.id == "Dict":
                if isinstance(node.slice, ast.Tuple):
                    return {"type": "object", "additionalProperties": _get_ast_type_info(node.slice.elts[1])}
            elif node.value.id == "Optional":
                return {"oneOf": [_get_ast_type_info(node.slice), {"type": "null"}]}
            elif node.value.id == "Literal":
                # 提取Literal所有值
                values = []
                if isinstance(node.slice, ast.Tuple):
                    for elt in node.slice.elts:
                        val = _get_constant_value(elt)
                        if val is not None:
                            values.append(val)
                else:
                    val = _get_constant_value(node.slice)
                    if val is not None:
                        values.append(val)

                # 推断类型
                if values:
                    inferred_type = type(values[0]).__name__
                else:
                    inferred_type = "string"

                return {"type": inferred_type, "enum": values}
    elif isinstance(node, ast.Constant):
        return {"type": type(node.value).__name__}
    elif isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub):
        # 处理负数
        value = _get_constant_value(node)
        return {"type": type(value).__name__}

    return {"type": "any"}


def _parse_function_signature(func) -> Dict[str, Any]:
    """
    使用 AST 解析函数签名

    Args:
        - func: 要解析的函数

    Returns:
        - dict: 函数签名信息
    """
    source = inspect.getsource(func)

    # 处理缩进问题
    # 1. 分割成行
    lines = source.splitlines()
    if not lines:
        return {}

    # 2. 找到第一行的缩进级别
    first_line = lines[0]
    indent = len(first_line) - len(first_line.lstrip())

    # 3. 对所有行去除相同的缩进
    normalized_lines = [line[indent:] if line.startswith(" " * indent) else line for line in lines]

    # 4. 重新组合成字符串
    normalized_source = "\n".join(normalized_lines)

    try:
        tree = ast.parse(normalized_source)
    except SyntaxError:
        return {}

    # 获取函数定义节点
    func_def = tree.body[0]
    if not isinstance(func_def, (ast.FunctionDef, ast.AsyncFunctionDef)):
        return {}

    parameters = {}
    required = []

    for arg in func_def.args.args:
        arg_name = getattr(arg, "arg", None)  # 使用 getattr 安全地获取参数名
        if not arg_name or arg_name == "self":
            continue

        param_info = {}
        # 获取类型注解
        if arg.annotation:
            param_info.update(_get_ast_type_info(arg.annotation))

        # 如果没有默认值，则为必需参数
        if len(func_def.args.defaults) < len(func_def.args.args) - func_def.args.args.index(arg):
            required.append(arg_name)

        parameters[arg_name] = param_info

    # 获取返回值类型
    returns = None
    if func_def.returns:
        returns = _get_ast_type_info(func_def.returns)

    return {"parameters": parameters, "required": required, "returns": returns}


def _parse_type_to_json_type(type_hint: Any) -> Dict[str, Any]:
    """
    将 Python 类型转换为 JSON Schema 类型

    Args:
        - type_hint: Python 类型提示

    Returns:
        - dict: JSON Schema 类型定义
    """
    if type_hint is str:
        return {"type": "string"}
    elif type_hint is int:
        return {"type": "integer"}
    elif type_hint is float:
        return {"type": "number"}
    elif type_hint is bool:
        return {"type": "boolean"}
    elif type_hint is list or str(type_hint).startswith("typing.List"):
        return {"type": "array"}
    elif type_hint is dict or str(type_hint).startswith("typing.Dict"):
        return {"type": "object"}
    elif isinstance(type_hint, type) and issubclass(type_hint, Enum):
        return {"type": "string", "enum": [e.value for e in type_hint]}
    elif hasattr(type_hint, "__origin__") and type_hint.__origin__ == Literal:
        return {"type": "string", "enum": list(type_hint.__args__)}
    else:
        # 对于其他类型，返回字符串形式的类型名
        return {"type": str(type_hint.__name__ if hasattr(type_hint, "__name__") else type_hint)}


def _parse_docstring_args(docstring: str) -> Tuple[Dict[str, Dict[str, Any]], List[str]]:
    """
    解析 docstring 中的参数说明

    Args:
        - docstring: 函数的 docstring

    Returns:
        - tuple: (参数列表, 必需参数列表)
    """
    # 匹配Args部分，使用更灵活的正则
    args_section = re.search(r"Args:(.*?)(?=Returns:|Raises:|\Z)", docstring, re.DOTALL)
    if not args_section:
        return {}, []

    args_text = args_section.group(1)
    lines = [line.rstrip() for line in args_text.splitlines() if line.strip()]

    properties = {}
    required = []
    current_param = None
    base_indent = None
    description = []

    for line in lines:
        # 计算缩进
        indent = len(line) - len(line.lstrip())
        stripped = line.strip()

        # 检测参数行（必须满足：包含冒号，且缩进等于基准缩进）
        param_match = re.match(r"-\s*(\w+)\s*(?:\(.*?\))?:\s*(.*)", stripped)
        if param_match:
            # 首次遇到参数行时设置基准缩进
            if base_indent is None:
                base_indent = indent

            # 只有缩进等于基准缩进才是新参数
            if indent == base_indent:
                # 保存前一个参数
                if current_param:
                    properties[current_param] = {"description": " ".join(description).strip()}
                    required.append(current_param)
                    description = []

                current_param = param_match.group(1)
                desc_part = param_match.group(2).strip()
                if desc_part:
                    description.append(desc_part)
                continue

        # 处理描述内容
        if current_param:
            # 子项处理（缩进大于基准缩进）
            if base_indent is not None and indent > base_indent:
                # 保留原始缩进结构
                clean_line = line[base_indent:].lstrip()
                if description and not description[-1].endswith("\n"):
                    description.append("\n" + clean_line)
                else:
                    description.append(clean_line)
            # 同级普通描述
            elif indent == base_indent:
                description.append(stripped)

    # 处理最后一个参数
    if current_param:
        properties[current_param] = {"description": " ".join(description).strip().replace("\n ", "\n")}
        required.append(current_param)

    return properties, required


def _parse_docstring_return(docstring: str) -> Optional[Dict[str, Any]]:
    """
    解析 docstring 中的返回值说明

    Args:
        - docstring: 函数的 docstring

    Returns:
        - Optional[Dict[str, Any]]: 返回值描述
    """
    returns_section = re.search(r"Returns:(.*?)(?:\n\s*[A-Za-z]+:|$)", docstring, re.DOTALL)
    if not returns_section:
        return None

    returns_text = returns_section.group(1).strip()

    # 尝试解析返回类型和描述
    type_desc_match = re.match(r"-?\s*([^:]+):\s*(.*)", returns_text)
    if type_desc_match:
        return_type = type_desc_match.group(1).strip()
        description = type_desc_match.group(2).strip()
    else:
        # 如果没有明确的类型标注，就把整个文本作为描述
        return_type = "any"
        description = returns_text

    return {"type": return_type.lower(), "description": description}


def _parse_docstring_description(docstring: str) -> str:
    """
    解析 docstring 中的描述部分

    Args:
        - docstring: 要解析的文档字符串

    Returns:
        - str: 描述文本，从开头到 Args: 或 Returns: 之前的部分
    """
    # 查找 Args: 或 Returns: 的位置
    args_pos = docstring.find("Args:")
    returns_pos = docstring.find("Returns:")

    # 找到最近的分隔位置
    end_pos = float("inf")
    for pos in [args_pos, returns_pos]:
        if pos != -1 and pos < end_pos:
            end_pos = pos

    # 如果没有找到分隔符，使用整个文档字符串
    if end_pos == float("inf"):
        description = docstring
    else:
        description = docstring[:end_pos]

    # 清理并返回描述文本
    return description.strip()


def parse_class_methods(cls: Type, exlude_methods: List[str] | None = None) -> Dict[str, Dict[str, Any]]:
    """
    解析类的所有公开方法，生成 JSON Schema 格式的文档

    Args:
        - cls: 要解析的类
        - exlude_methods: 要排除的方法（该方法不会被处理）

    Returns:
        - Dict[str, Dict[str, Any]]: 方法文档字典
    """
    result = {}

    for name, method in inspect.getmembers(cls, predicate=inspect.isfunction):
        # 跳过内部方法
        if name.startswith("_"):
            continue

        if exlude_methods and name in exlude_methods:
            continue

        # 获取方法的文档字符串
        docstring = inspect.getdoc(method)
        if not docstring:
            continue

        # 使用 AST 解析函数签名（这是主要的参数来源）
        signature_info = _parse_function_signature(method)

        # 获取类型提示作为备用
        try:
            type_hints = get_type_hints(method)
        except Exception:
            type_hints = {}

        # 解析 docstring 中的参数说明，仅用于补充描述
        doc_properties, _ = _parse_docstring_args(docstring)

        # 以签名参数为准，只从 docstring 获取描述信息
        properties = signature_info.get("parameters", {}).copy()
        for param_name, param_info in properties.items():
            if param_name in doc_properties:
                param_info["description"] = doc_properties[param_name]["description"]

        # 必需参数只从签名中获取
        required = signature_info.get("required", [])

        # 解析返回值
        returns = _parse_docstring_return(docstring)
        if signature_info.get("returns"):
            if returns:
                # 保留签名中的类型信息，只使用 docstring 中的描述
                returns_type = signature_info["returns"].copy()
                if "description" in returns:
                    returns_type["description"] = returns["description"]
                returns = returns_type
            else:
                returns = signature_info["returns"]
        elif "return" in type_hints:
            returns = _parse_type_to_json_type(type_hints["return"])

        # 构建方法文档
        method_doc = {
            "description": _parse_docstring_description(docstring),  # 第一行作为简要描述
            "parameters": {"type": "object", "properties": properties, "required": required},
        }

        if returns:
            method_doc["return"] = returns

        result[name] = method_doc

    return result


if __name__ == "__main__":
    import json

    from mozi_ai_x.simulation.active_unit.aircraft import CAircraft

    class_methods = parse_class_methods(CAircraft, ["parse"])

    with open("CAircraft.json", "w", encoding="utf-8") as f:
        json.dump(class_methods, f, indent=2, ensure_ascii=False)
